import json
import torch
from torch.utils.data import Dataset

def load_data(max_len):
    # 文件路径
    train_file_path = "data/KUAKE-QIC_train.json"
    dev_file_path = "data/KUAKE-QIC_dev.json"
    # 打开文件
    with open(train_file_path, 'r', encoding='utf-8') as file:
        train_data = json.load(file)

    with open(dev_file_path, 'r', encoding='utf-8') as file:
        dev_data = json.load(file)
    # 加载标签名列表
    class_names = open("data/KUAKE-QIC_class.txt", 'r', encoding='utf-8').read().split('\n')
    # 从文件中提取出文本内容和标签内容
    train_text = []
    train_label = []
    for item in train_data:
        text = item['query']
        label = item['label']
        train_text.append(text)
        train_label.append(class_names.index(label))
    dev_text = []
    dev_label = []
    for item in dev_data:
        text = item['query']
        label = item['label']
        dev_text.append(text)
        dev_label.append(class_names.index(label))
    # 裁剪或填充至你设置好的最大长度
    train_text = [truncate_pad(text, max_len) for text in train_text]
    dev_text = [truncate_pad(text, max_len) for text in dev_text]
    # 构建词汇表
    vocab = set()
    for text in train_text + dev_text:
        for char in text:
            vocab.add(char)
    vocab = list(vocab)
    vocab.append('<pad>')
    vocab.append('<unk>')
    # 把词表保存下来
    with open("vocab.txt", 'w', encoding='utf-8') as file:
        file.write('\n'.join(vocab))
    # 转换为索引
    train_text = [[vocab.index(char) if char in vocab else vocab.index('<unk>') for char in text] for text in train_text]
    dev_text = [[vocab.index(char) if char in vocab else vocab.index('<unk>') for char in text] for text in dev_text]
    return train_text, train_label, dev_text, dev_label, vocab

# 裁剪或填充
def truncate_pad(text, max_len):
    text = list(text)
    if len(text) > max_len:
        text = text[:max_len]
    else:
        text = text + ['<pad>'] * (max_len - len(text))
    return text

class KuakeDataset(Dataset):
    def __init__(self, text, label):
        self.text = text
        self.label = label
    def __len__(self):
        return len(self.text)
    def __getitem__(self, idx):
        text = self.text[idx]
        label = self.label[idx]
        item = {'input_ids': torch.tensor(text), 'labels': torch.tensor(label)}
        return item

def load_dataset(max_len):
    train_text, train_label, dev_text, dev_label, vocab = load_data(max_len)
    train_dataset = KuakeDataset(train_text, train_label)
    dev_dataset = KuakeDataset(dev_text, dev_label)
    return train_dataset, dev_dataset, vocab

